# 读取数据
import numpy as np
import pandas as pd
import torch
from matplotlib import pyplot as plt
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch import nn
from sklearn.metrics import mean_absolute_error as mae, r2_score
import torch.optim as optim

data = pd.read_csv('./data/ETTh1.csv')

# 提取特征和标签
f1 = data['HUFL'].values
f2 = data['HULL'].values
f3 = data['MUFL'].values
f4 = data['MULL'].values
OT = data['OT'].values


# 定义历史和未来的窗口大小
history_size = 10
future_size = 1
# 那么采样步长设为2，避免预测值重叠
step_size = 2


features = []
labels = []
for i in range(history_size, len(f1) - 2, step_size):  # 按步长取样（不重叠）
    history_f1 = f1[i - history_size:i]
    history_f2 = f2[i - history_size:i]
    history_f3 = f3[i - history_size:i]
    history_f4 = f4[i - history_size:i] #shape: (10, )

    # 输入特征
    feature = np.stack((history_f1, history_f2, history_f3, history_f4), axis = 1) #维度：(10, 4)
    features.append(feature)

    # 输出标签
    label = OT[i:i + future_size] #未来feature_size的标签
    labels.append(label)

# 转换为 NumPy 数组
features = np.array(features) #list -> ndarray， 维度为（8704，10，4）
labels = np.array(labels)

# 归一化，对 features 和 labels 数据进行标准化处理，使它们的均值为 0，标准差为 1。
scaler_x = StandardScaler()
scaler_y = StandardScaler()

### features_nomalized(87040,4)
features_nomalized = scaler_x.fit_transform(features.reshape(-1, 4)) #转换成二维数组进行归一化（归一化只能对二维数组进行）
labels_nomalized = scaler_y.fit_transform(labels.reshape(-1, 1)) #这里的1表示只有一列标签

features = features_nomalized.reshape(-1, history_size, 4) #将转换后的二维数组还原为三维数组
labels = labels_nomalized.reshape(-1, 1)


# 划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.1)


# 将特征转换为三维张量，形状为 [样本数, 时间序列长度, 特征数] -> [xx, history_size, 4]
x_train_tensor = torch.tensor(x_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
x_test_tensor = torch.tensor(x_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32)


# 2. 创建RNN模型
class RNNModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(RNNModel, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True) #输入层特征个数，隐藏层特征个数，隐藏层个数
        self.fc = nn.Linear(hidden_size, output_size) #隐藏层特征个数，输出层特征个数

    def forward(self, x):
        h0 = torch.zeros(self.rnn.num_layers, x.size(0), self.rnn.hidden_size).to(x.device)
        '''
            创建一个形状为 (num_layers, batch_size, hidden_size) 的张量
            num_layers为rnn的层数，也就是隐藏层的个数
        '''
        out, _ = self.rnn(x, h0)
        return self.fc(out[:, -1, :])  # 只取最后一个时间步的输出


# 实例化模型
input_size = 4  # 使用特征的列数作为输入特征数
hidden_size = 64  # 隐藏层大小
num_layers = 2  # RNN层数
output_size = 1  # 输出特征数
model = RNNModel(input_size, hidden_size, num_layers, output_size)

# 3. 设置损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练模型
num_epochs = 1000
for epoch in range(num_epochs):

    model.train()  # 将模型调至训练模式
    # 前向传播
    outputs = model(x_train_tensor)
    loss = criterion(outputs, y_train_tensor)

    # 后向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 5. 预测
model.eval()
with torch.no_grad():
    y_pred_tensor = model(x_test_tensor)

y_pred = scaler_y.inverse_transform(y_pred_tensor.numpy())  # 将预测值逆归一化
y_test = scaler_y.inverse_transform(y_test_tensor.numpy())  # 逆归一化真实值

# 评估指标
r2 = r2_score(y_test, y_pred, multioutput='uniform_average')
mae_score = mae(y_test, y_pred)
print(f"R^2 score: {r2:.4f}")
print(f"MAE: {mae_score:.4f}")

# 绘制对比图, 配置中文字体支持以及调整符号的显示方式
plt.rcParams['font.sans-serif'] = ['SimSun']
plt.rcParams['axes.unicode_minus'] = False

plt.figure(figsize=(10, 6))  # 创建一个新的 Matplotlib 图形，并设置图形的宽高尺寸
sample_indices = np.arange(len(y_test))  # 用于创建一个[0, len(y_test)]的等差数列的数组
plt.plot(sample_indices, y_test, color='blue', alpha=0.5, label='真实值')  # 绘制折线图，sample_indices为x轴数据，y_test为y轴数据
plt.plot(sample_indices, y_pred, color='red', alpha=0.5, label='预测值')
plt.xlabel('样本索引')
plt.ylabel('OT值')
plt.title('OT的实际值与预测值对比图')
plt.legend(['实际数据', '预测数据'])  # 显示图例
plt.grid(True)  # 显示网格线
plt.show()
